import sys
from os import path

sys.path.append(path.join(path.dirname(__file__), ".."))
sys.path.append(path.join(path.dirname(__file__), "../.."))
import gym
from dqnv2.DQN import Deep_Q_Network


def main():
    step = 0
    for episode in range(300):
        # initial observation
        observation = env.reset()

        while True:
            # fresh env
            env.render()

            # RL choose action based on observation
            action = RL.choose_action(observation)

            # RL take action and get next observation and reward
            observation_, reward, done, truncated, info = env.step(action)

            RL.store_transition(observation, action, reward, done, observation_)

            if (step > 200) and (step % 5 == 0):
                RL.learn()
                RL.soft_update(1)

            # swap observation
            observation = observation_

            # break while loop when end of this episode
            if done:
                break
            step += 1
        print("episode %d" % (episode + 1))
    # end of game
    print("game over")
    env.close()


if __name__ == "__main__":
    env = gym.make("CartPole-v1")
    RL = Deep_Q_Network(
        n_actions=env.action_space.n, n_features=env.observation_space.shape[0]
    )
    main()
